"""
Feature segmentation system for CodeMCP.

Provides user segment targeting and evaluation logic for feature flags
with support for complex targeting rules and user attribute matching.
"""

import hashlib
import re
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Set, Union, Callable
from dataclasses import dataclass, field
from enum import Enum

from ...core.error_handler import CodeMCPError, log_info, log_warning, log_debug, handle_error


class SegmentOperator(Enum):
    """Segment targeting operators."""
    EQUALS = "equals"
    NOT_EQUALS = "not_equals"
    CONTAINS = "contains"
    NOT_CONTAINS = "not_contains"
    STARTS_WITH = "starts_with"
    ENDS_WITH = "ends_with"
    REGEX_MATCH = "regex_match"
    GREATER_THAN = "greater_than"
    LESS_THAN = "less_than"
    GREATER_EQUAL = "greater_equal"
    LESS_EQUAL = "less_equal"
    IN = "in"
    NOT_IN = "not_in"
    IS_NULL = "is_null"
    IS_NOT_NULL = "is_not_null"


class SegmentLogic(Enum):
    """Segment logic operators."""
    AND = "and"
    OR = "or"
    NOT = "not"


@dataclass
class SegmentRule:
    """Individual segment targeting rule."""
    attribute: str
    operator: SegmentOperator
    value: Any = None
    case_sensitive: bool = True
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert segment rule to dictionary."""
        return {
            "attribute": self.attribute,
            "operator": self.operator.value,
            "value": self.value,
            "case_sensitive": self.case_sensitive
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SegmentRule":
        """Create segment rule from dictionary."""
        return cls(
            attribute=data["attribute"],
            operator=SegmentOperator(data["operator"]),
            value=data.get("value"),
            case_sensitive=data.get("case_sensitive", True)
        )


@dataclass
class SegmentGroup:
    """Group of segment rules with logic operator."""
    logic: SegmentLogic
    rules: List[SegmentRule] = field(default_factory=list)
    groups: List["SegmentGroup"] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert segment group to dictionary."""
        return {
            "logic": self.logic.value,
            "rules": [rule.to_dict() for rule in self.rules],
            "groups": [group.to_dict() for group in self.groups]
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SegmentGroup":
        """Create segment group from dictionary."""
        return cls(
            logic=SegmentLogic(data["logic"]),
            rules=[SegmentRule.from_dict(rule_data) for rule_data in data.get("rules", [])],
            groups=[SegmentGroup.from_dict(group_data) for group_data in data.get("groups", [])]
        )


@dataclass
class UserSegment:
    """User segment definition."""
    name: str
    description: str
    targeting_group: SegmentGroup
    priority: int = 0
    active: bool = True
    environment: str = "all"
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.utcnow)
    updated_at: datetime = field(default_factory=datetime.utcnow)
    version: int = 1
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert user segment to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "targeting_group": self.targeting_group.to_dict(),
            "priority": self.priority,
            "active": self.active,
            "environment": self.environment,
            "metadata": self.metadata,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "version": self.version
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "UserSegment":
        """Create user segment from dictionary."""
        return cls(
            name=data["name"],
            description=data["description"],
            targeting_group=SegmentGroup.from_dict(data["targeting_group"]),
            priority=data.get("priority", 0),
            active=data.get("active", True),
            environment=data.get("environment", "all"),
            metadata=data.get("metadata", {}),
            created_at=datetime.fromisoformat(data.get("created_at", datetime.utcnow().isoformat())),
            updated_at=datetime.fromisoformat(data.get("updated_at", datetime.utcnow().isoformat())),
            version=data.get("version", 1)
        )


class SegmentationError(CodeMCPError):
    """Segmentation specific errors."""
    
    def __init__(self, message: str, segment_name: str = None, **kwargs):
        super().__init__(message, "SEGMENTATION_ERROR", **kwargs)
        if segment_name:
            self.details["segment_name"] = segment_name


class FeatureSegmentation:
    """
    Feature segmentation system for CodeMCP.
    
    Provides advanced user segment targeting and evaluation logic
    with support for complex rules and user attribute matching.
    """
    
    def __init__(self):
        # Segment registry
        self.segments: Dict[str, UserSegment] = {}
        self.segment_cache: Dict[str, Dict[str, Any]] = {}
        
        # Performance tracking
        self.evaluation_stats = {
            "total_evaluations": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "avg_evaluation_time": 0.0,
            "rule_evaluation_count": {},
            "segment_match_count": {}
        }
        
        # Configuration
        self.cache_ttl_seconds = 300  # 5 minutes
        self.max_cache_size = 5000
        self.enable_regex_cache = True
        self.regex_cache: Dict[str, re.Pattern] = {}
        
        log_info("Feature segmentation system initialized")
    
    async def initialize(self):
        """Initialize the segmentation system."""
        try:
            # Load default segments
            await self._load_default_segments()
            
            log_info("Feature segmentation system initialized successfully")
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "segmentation_init"})
            raise SegmentationError(f"Failed to initialize segmentation system: {error_response['message']}")
    
    async def _load_default_segments(self):
        """Load default user segments for CodeMCP GSI Wiki Generator."""
        try:
            # Developer segment
            developer_segment = UserSegment(
                name="developers",
                description="Software developers and engineers",
                targeting_group=SegmentGroup(
                    logic=SegmentLogic.OR,
                    rules=[
                        SegmentRule(
                            attribute="user_groups",
                            operator=SegmentOperator.CONTAINS,
                            value="developers"
                        ),
                        SegmentRule(
                            attribute="user_groups",
                            operator=SegmentOperator.CONTAINS,
                            value="engineers"
                        ),
                        SegmentRule(
                            attribute="role",
                            operator=SegmentOperator.IN,
                            value=["developer", "engineer", "tech_lead"]
                        )
                    ]
                ),
                priority=10
            )
            
            # Beta users segment
            beta_users_segment = UserSegment(
                name="beta_users",
                description="Beta testing users",
                targeting_group=SegmentGroup(
                    logic=SegmentLogic.AND,
                    rules=[
                        SegmentRule(
                            attribute="user_groups",
                            operator=SegmentOperator.CONTAINS,
                            value="beta_users"
                        ),
                        SegmentRule(
                            attribute="beta_opt_in",
                            operator=SegmentOperator.EQUALS,
                            value=True
                        )
                    ]
                ),
                priority=20
            )
            
            # GSI researchers segment
            gsi_researchers_segment = UserSegment(
                name="gsi_researchers",
                description="GSI/DRP-4DVar researchers and scientists",
                targeting_group=SegmentGroup(
                    logic=SegmentLogic.OR,
                    rules=[
                        SegmentRule(
                            attribute="department",
                            operator=SegmentOperator.CONTAINS,
                            value="atmospheric"
                        ),
                        SegmentRule(
                            attribute="research_area",
                            operator=SegmentOperator.IN,
                            value=["data_assimilation", "gsi", "drp4dvar", "ensemble"]
                        ),
                        SegmentRule(
                            attribute="user_groups",
                            operator=SegmentOperator.CONTAINS,
                            value="gsi_researchers"
                        )
                    ]
                ),
                priority=15
            )
            
            # Power users segment
            power_users_segment = UserSegment(
                name="power_users",
                description="Power users with advanced usage patterns",
                targeting_group=SegmentGroup(
                    logic=SegmentLogic.AND,
                    rules=[
                        SegmentRule(
                            attribute="usage_frequency",
                            operator=SegmentOperator.GREATER_THAN,
                            value=100  # 100+ sessions
                        ),
                        SegmentRule(
                            attribute="advanced_features_used",
                            operator=SegmentOperator.GREATER_EQUAL,
                            value=5
                        )
                    ]
                ),
                priority=25
            )
            
            # Production environment segment
            production_segment = UserSegment(
                name="production_environment",
                description="Production environment users",
                targeting_group=SegmentGroup(
                    logic=SegmentLogic.AND,
                    rules=[
                        SegmentRule(
                            attribute="environment",
                            operator=SegmentOperator.EQUALS,
                            value="production"
                        )
                    ]
                ),
                priority=5,
                environment="production"
            )
            
            # Add segments
            segments = {
                "developers": developer_segment,
                "beta_users": beta_users_segment,
                "gsi_researchers": gsi_researchers_segment,
                "power_users": power_users_segment,
                "production_environment": production_segment
            }
            
            self.segments.update(segments)
            log_info(f"Loaded {len(segments)} default user segments")
            
        except Exception as e:
            log_warning(f"Failed to load default segments: {e}")
    
    async def register_segment(self, segment: UserSegment) -> bool:
        """Register a new user segment."""
        try:
            if segment.name in self.segments:
                existing_segment = self.segments[segment.name]
                segment.version = existing_segment.version + 1
                segment.updated_at = datetime.utcnow()
                log_info(f"Updated existing segment: {segment.name} (v{segment.version})")
            else:
                log_info(f"Registered new segment: {segment.name}")
            
            self.segments[segment.name] = segment
            
            # Clear cache for this segment
            await self._clear_segment_cache(segment.name)
            
            return True
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "register_segment", "segment_name": segment.name})
            raise SegmentationError(f"Failed to register segment: {error_response['message']}", segment.name)
    
    async def unregister_segment(self, segment_name: str) -> bool:
        """Unregister a user segment."""
        try:
            if segment_name not in self.segments:
                log_warning(f"Segment not found: {segment_name}")
                return False
            
            del self.segments[segment_name]
            await self._clear_segment_cache(segment_name)
            
            log_info(f"Unregistered segment: {segment_name}")
            return True
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "unregister_segment", "segment_name": segment_name})
            raise SegmentationError(f"Failed to unregister segment: {error_response['message']}", segment_name)
    
    async def evaluate_user_segments(self, user_context: Dict[str, Any]) -> Set[str]:
        """Evaluate which segments a user belongs to."""
        import time
        start_time = time.time()
        
        try:
            # Generate cache key
            cache_key = self._generate_user_cache_key(user_context)
            
            # Check cache
            cached_result = await self._get_cached_segments(cache_key)
            if cached_result is not None:
                self.evaluation_stats["cache_hits"] += 1
                return cached_result
            
            self.evaluation_stats["cache_misses"] += 1
            
            # Evaluate segments
            matching_segments = set()
            
            # Sort segments by priority (higher priority first)
            sorted_segments = sorted(
                self.segments.values(),
                key=lambda s: s.priority,
                reverse=True
            )
            
            for segment in sorted_segments:
                if not segment.active:
                    continue
                
                # Check environment match
                if segment.environment != "all" and user_context.get("environment") != segment.environment:
                    continue
                
                # Evaluate targeting group
                if await self._evaluate_segment_group(segment.targeting_group, user_context):
                    matching_segments.add(segment.name)
                    
                    # Update segment match statistics
                    self.evaluation_stats["segment_match_count"][segment.name] = (
                        self.evaluation_stats["segment_match_count"].get(segment.name, 0) + 1
                    )
            
            # Cache the result
            await self._cache_user_segments(cache_key, matching_segments, time.time() - start_time)
            
            return matching_segments
            
        except Exception as e:
            error_response = handle_error(e, {
                "operation": "evaluate_user_segments",
                "user_context": user_context
            })
            log_warning(f"Segment evaluation failed: {error_response['message']}")
            return set()
        
        finally:
            # Update stats
            self.evaluation_stats["total_evaluations"] += 1
    
    async def _evaluate_segment_group(self, group: SegmentGroup, user_context: Dict[str, Any]) -> bool:
        """Evaluate a segment group against user context."""
        try:
            if group.logic == SegmentLogic.AND:
                # All rules and groups must match
                for rule in group.rules:
                    if not await self._evaluate_segment_rule(rule, user_context):
                        return False
                
                for nested_group in group.groups:
                    if not await self._evaluate_segment_group(nested_group, user_context):
                        return False
                
                return True
            
            elif group.logic == SegmentLogic.OR:
                # At least one rule or group must match
                for rule in group.rules:
                    if await self._evaluate_segment_rule(rule, user_context):
                        return True
                
                for nested_group in group.groups:
                    if await self._evaluate_segment_group(nested_group, user_context):
                        return True
                
                return False
            
            elif group.logic == SegmentLogic.NOT:
                # None of the rules or groups should match
                for rule in group.rules:
                    if await self._evaluate_segment_rule(rule, user_context):
                        return False
                
                for nested_group in group.groups:
                    if await self._evaluate_segment_group(nested_group, user_context):
                        return False
                
                return True
            
            return False
            
        except Exception as e:
            log_debug(f"Failed to evaluate segment group: {e}")
            return False
    
    async def _evaluate_segment_rule(self, rule: SegmentRule, user_context: Dict[str, Any]) -> bool:
        """Evaluate a single segment rule against user context."""
        try:
            # Update rule evaluation statistics
            self.evaluation_stats["rule_evaluation_count"][rule.attribute] = (
                self.evaluation_stats["rule_evaluation_count"].get(rule.attribute, 0) + 1
            )
            
            # Get attribute value from user context
            attribute_value = user_context.get(rule.attribute)
            
            # Handle null checks first
            if rule.operator == SegmentOperator.IS_NULL:
                return attribute_value is None
            
            if rule.operator == SegmentOperator.IS_NOT_NULL:
                return attribute_value is not None
            
            # If attribute doesn't exist and it's not a null check, return False
            if attribute_value is None:
                return False
            
            # Convert values for case-insensitive comparison
            if not rule.case_sensitive and isinstance(attribute_value, str) and isinstance(rule.value, str):
                attribute_value = attribute_value.lower()
                comparison_value = rule.value.lower()
            else:
                comparison_value = rule.value
            
            # Evaluate based on operator
            if rule.operator == SegmentOperator.EQUALS:
                return attribute_value == comparison_value
            
            elif rule.operator == SegmentOperator.NOT_EQUALS:
                return attribute_value != comparison_value
            
            elif rule.operator == SegmentOperator.CONTAINS:
                if isinstance(attribute_value, str):
                    return comparison_value in attribute_value
                elif isinstance(attribute_value, (list, set)):
                    return comparison_value in attribute_value
                return False
            
            elif rule.operator == SegmentOperator.NOT_CONTAINS:
                if isinstance(attribute_value, str):
                    return comparison_value not in attribute_value
                elif isinstance(attribute_value, (list, set)):
                    return comparison_value not in attribute_value
                return True
            
            elif rule.operator == SegmentOperator.STARTS_WITH:
                if isinstance(attribute_value, str):
                    return attribute_value.startswith(comparison_value)
                return False
            
            elif rule.operator == SegmentOperator.ENDS_WITH:
                if isinstance(attribute_value, str):
                    return attribute_value.endswith(comparison_value)
                return False
            
            elif rule.operator == SegmentOperator.REGEX_MATCH:
                if isinstance(attribute_value, str):
                    pattern = self._get_regex_pattern(comparison_value, rule.case_sensitive)
                    return bool(pattern.search(attribute_value))
                return False
            
            elif rule.operator == SegmentOperator.GREATER_THAN:
                return float(attribute_value) > float(comparison_value)
            
            elif rule.operator == SegmentOperator.LESS_THAN:
                return float(attribute_value) < float(comparison_value)
            
            elif rule.operator == SegmentOperator.GREATER_EQUAL:
                return float(attribute_value) >= float(comparison_value)
            
            elif rule.operator == SegmentOperator.LESS_EQUAL:
                return float(attribute_value) <= float(comparison_value)
            
            elif rule.operator == SegmentOperator.IN:
                if isinstance(comparison_value, (list, set, tuple)):
                    return attribute_value in comparison_value
                return False
            
            elif rule.operator == SegmentOperator.NOT_IN:
                if isinstance(comparison_value, (list, set, tuple)):
                    return attribute_value not in comparison_value
                return True
            
            return False
            
        except Exception as e:
            log_debug(f"Failed to evaluate segment rule {rule.attribute}: {e}")
            return False
    
    def _get_regex_pattern(self, pattern: str, case_sensitive: bool) -> re.Pattern:
        """Get compiled regex pattern with caching."""
        cache_key = f"{pattern}:{case_sensitive}"
        
        if cache_key not in self.regex_cache:
            flags = 0 if case_sensitive else re.IGNORECASE
            self.regex_cache[cache_key] = re.compile(pattern, flags)
        
        return self.regex_cache[cache_key]
    
    def _generate_user_cache_key(self, user_context: Dict[str, Any]) -> str:
        """Generate cache key for user context."""
        # Create a stable string representation of the context
        context_items = sorted(user_context.items())
        context_str = str(context_items)
        
        # Hash for compact key
        return hashlib.sha256(context_str.encode()).hexdigest()[:16]
    
    async def _get_cached_segments(self, cache_key: str) -> Optional[Set[str]]:
        """Get cached segment evaluation result."""
        try:
            if cache_key in self.segment_cache:
                cache_entry = self.segment_cache[cache_key]
                if time.time() < cache_entry["expires_at"]:
                    return cache_entry["segments"]
                else:
                    # Remove expired entry
                    del self.segment_cache[cache_key]
            
            return None
            
        except Exception as e:
            log_debug(f"Failed to get cached segments: {e}")
            return None
    
    async def _cache_user_segments(self, cache_key: str, segments: Set[str], evaluation_time: float):
        """Cache user segment evaluation result."""
        try:
            # Clean cache if it's too large
            if len(self.segment_cache) >= self.max_cache_size:
                await self._clean_segment_cache()
            
            import time
            expires_at = time.time() + self.cache_ttl_seconds
            self.segment_cache[cache_key] = {
                "segments": segments,
                "evaluation_time": evaluation_time,
                "expires_at": expires_at
            }
            
        except Exception as e:
            log_debug(f"Failed to cache user segments: {e}")
    
    async def _clean_segment_cache(self):
        """Clean expired cache entries."""
        try:
            import time
            current_time = time.time()
            expired_keys = [
                key for key, entry in self.segment_cache.items()
                if current_time >= entry["expires_at"]
            ]
            
            for key in expired_keys:
                del self.segment_cache[key]
            
            # If still too large, remove oldest entries
            if len(self.segment_cache) >= self.max_cache_size:
                sorted_entries = sorted(
                    self.segment_cache.items(),
                    key=lambda x: x[1]["expires_at"]
                )
                
                # Remove oldest 25%
                num_to_remove = len(sorted_entries) // 4
                for key, _ in sorted_entries[:num_to_remove]:
                    del self.segment_cache[key]
            
        except Exception as e:
            log_warning(f"Failed to clean segment cache: {e}")
    
    async def _clear_segment_cache(self, segment_name: str):
        """Clear cache entries that might be affected by segment changes."""
        # For simplicity, clear entire cache when a segment changes
        # In production, you might want to be more selective
        self.segment_cache.clear()
    
    async def get_segment(self, segment_name: str) -> Optional[UserSegment]:
        """Get a user segment by name."""
        return self.segments.get(segment_name)
    
    async def list_segments(self, environment: str = None, active: bool = None) -> List[UserSegment]:
        """List user segments with optional filtering."""
        segments = list(self.segments.values())
        
        if environment:
            segments = [s for s in segments if s.environment == environment or s.environment == "all"]
        
        if active is not None:
            segments = [s for s in segments if s.active == active]
        
        return sorted(segments, key=lambda s: s.priority, reverse=True)
    
    async def get_segmentation_stats(self) -> Dict[str, Any]:
        """Get segmentation performance statistics."""
        stats = self.evaluation_stats.copy()
        stats["cache_size"] = len(self.segment_cache)
        stats["regex_cache_size"] = len(self.regex_cache)
        stats["segments_registered"] = len(self.segments)
        stats["cache_hit_rate"] = (
            self.evaluation_stats["cache_hits"] / 
            max(1, self.evaluation_stats["cache_hits"] + self.evaluation_stats["cache_misses"])
        )
        
        return stats